#include<iostream>
#include<cstdlib>
#include<ctime>
#include<cstring>
#include<assert.h>
#include<cv.h>
#include<highgui.h>

#define M 5000000 //max number of pixels
#define D 3 // max number of chanels
#define K_MAX 1024 // number of clusters
#define C 256 // max intensity color 

//blusters
int center[K_MAX][D];
int ccount[K_MAX];
int K;

//blocks
int bl_iter[C][C][C], iter;
int bl_count[C][C][C];
int bl_point[C][C][C][D];
int bl_clust[C][C][C];
int bl_used[C*C*C][D];
int n_used;

inline
int sqr(int x)
{
	return x*x;
}


int dist2(int* a,int* b)
{
	int ans=0;
	for(int i=0;i<D;++i)
		ans += sqr(a[i]-b[i]);
	return ans;
}
/**
* R : color division factor
*/
inline 
void find_clust(int R)
{
	int bl_center[D];

	for(int bl = 0; bl < n_used; ++bl)
	{
		int r = bl_used[bl][0], g = bl_used[bl][1], b = bl_used[bl][2];
	
		for(int d=0;d<D;++d)
			bl_center[d] = bl_point[r][g][b][d]/bl_count[r][g][b];

		int bd = INT_MAX, cl;
		//find the closest cluster
		for(int k=0;k<K;++k)
		{
			int qd = dist2(bl_center, center[k]);
			if(qd < bd)
			{
				bd = qd;
				cl = k;
			}
		}
		bl_clust[r][g][b] = cl;
	}
}
/** Realocates the cluster centers
* precondition : cl_clust updateded (call find clust)
* R : color division factor
*/
inline
void calc_clust(int R)
{
	memset(center, 0, sizeof center);
	memset(ccount, 0, sizeof ccount);
	
	for(int bl = 0; bl < n_used; ++bl)
	{
		int r = bl_used[bl][0], g = bl_used[bl][1], b = bl_used[bl][2];

		int cl = bl_clust[r][g][b];
		
		ccount[cl] += bl_count[r][g][b];
		for(int d=0;d<D;++d)
		{
			center[cl][d] += bl_point[r][g][b][d];
		}
	}

	for(int k=0; k < K; ++k)
		for(int d=0; d<D; ++d)
			if( ccount[k] )
			{
				center[k][d] /= ccount[k];
			}else
			{
				int bl = rand()%n_used;
				
				int r = bl_used[bl][0], g = bl_used[bl][1], b = bl_used[bl][2];

				for(int d=0;d<D;++d)
					center[k][d] = bl_point[r][g][b][d]/bl_count[r][g][b];
			}
}
/*
* R : color division factor
* I : number of iterations
*/
inline
void k_means(int R, int I)
{	
	while(I--)
	{
		calc_clust(R); //O(N);
		find_clust(R); //O(N*K);
	}
}


inline
void find_parameters(int px, int& I, int& R)
{
	// first r to satisfy
	R = 0;
	I=1;
	
	for(;R<=8;++R)
	{
		if( (1LL<<(24-3*R))*(K*I) <= px )
			break;
	}
	
	for(++I;;++I)
	{
		if( (1LL<<(24-3*R))*(K*I) > px )
			break;
	}
	--I;
	
	//R =3, I =2;

}

int main(int argc, char* argv[])
{
	srand(time(0));
	if(argc < 3 )
	{
		printf("Command error: ./compress [file_to_be_compressed] [K] ([times])?\n");
		return -1;
	}

	if( sscanf(argv[2],"%d", &K) !=1 )
	{
		printf("Second parameter K must be a number\n");
		return -1;
	}
	if( K > K_MAX )
	{
		printf("K must no be bigger than %d", K_MAX);
		return -1;
	}
	
	int T ; // used to define more video processing
	if(argc != 4 ||   sscanf(argv[3],"%d", &T) !=1 )
	{
		T =1;
	}

	cvNamedWindow("Compression",CV_WINDOW_AUTOSIZE);
	cvNamedWindow("Original",CV_WINDOW_AUTOSIZE);
	CvCapture* capture = cvCreateFileCapture(argv[1]);
	
	if( !capture )
	{
		printf( "\"%s\": File not found\n",argv[1]);
		return -1;
	}


	IplImage* frame;
	
	int iter = 0;
	bool fst = 1;
	// random initial centers
	for(int k=0; k<K; ++k)
		for(int d=0; d<D; ++d)
		{
			center[k][d] = rand() & 255;
		}

	// compress frame
	while(1)
	{
		// global variables
		++iter; 
		n_used = 0; 
		///

		frame = cvQueryFrame(capture);
		if(!frame)
			break;
		
		cvShowImage("Original",frame);
		
		//find parameters
		int n = frame->height ;// rows
		int m = frame->width;//columns
		int I,R;
		find_parameters(T*n*m, I, R);
		//
		if(fst)
		{
			printf("K = %d, T= %d, R = %d, I = %d\n",K, T, R, I);
			fst = 0;
		}

		int ch = frame->nChannels;
		assert(ch == D);

		//calc blockx
		int co[D];
		for(int r=0, i=0;r<n;++r)
			for(int c=0;c<m;++c,++i)
			{
				for(int j=0;j<ch;++j)
				{
					co[j] = (unsigned char)frame->imageData[ch*(r*m+c)+j]; // singed values
				}
				// update region
				int x= co[0] >> R, y = co[1] >> R, z= co[2] >> R;
				 
				if( bl_iter[x][y][z] != iter)
				{
					bl_iter[x][y][z] = iter;
					bl_used[n_used][0] = x;
					bl_used[n_used][1] = y;
					bl_used[n_used][2] = z;

					++n_used;
					
					bl_count[x][y][z] = 0; 
					for(int d=0;d<D;++d)
						bl_point[x][y][z][d] = 0;
				}

				++bl_count[x][y][z];

				for(int j=0;j<ch;++j)
					bl_point[x][y][z][j] += co[j];
			}
		
		// update the centers with the definied blocks
		find_clust(R);
		k_means(R, I);

		for(int r=0, i=0;r<n;++r)
			for(int c=0;c<m;++c,++i)
			{
				//find color
				for(int j=0;j<ch;++j)
					co[j] = (unsigned char) frame->imageData[ch*(r*m+c)+j];
				//find cluster
				int cl = bl_clust[co[0]>>R][co[1]>>R][co[2]>>R];

				for(int j=0;j<ch;++j)
				{
					frame->imageData[ch*(r*m+c)+j] = (unsigned char)center[cl][j] ;
				}
			}

		cvShowImage("Compression",frame);

		++iter;
		//printf("Iter : %d\n",iter);
		char c = cvWaitKey(33);
		if(c == 27)
			break;
		//printf("\n");
	}
	cvReleaseCapture(&capture);
	cvDestroyWindow("Compression");
	return 0;
}
